//! Nuclei核心功能模块
//! 提供模板解析、变量渲染和响应匹配等核心功能

use anyhow::{anyhow, Result};
use log::{debug, error, info, warn};
use regex::Regex;
use reqwest::header::{HeaderMap, HeaderName, HeaderValue};
use serde::{Deserialize, Serialize};
use serde_json::{Value as JsonValue};
use serde_yaml::Value as YamlValue;
use std::collections::HashMap;
use std::fs;
use std::path::Path;
use std::time::Duration;
use rand::Rng;

/// Nuclei请求结构定义
/// 表示模板中的单个HTTP请求配置
#[derive(Debug, Deserialize, Clone, Serialize)]
pub struct NucleiRequest {
    /// HTTP请求方法，如GET、POST等
    #[serde(default)]
    pub method: String,
    /// 请求路径，支持变量替换
    #[serde(default)]
    pub path: String,
    /// 请求头配置
    #[serde(default)]
    pub headers: HashMap<String, String>,
    /// 请求体内容
    #[serde(default)]
    pub body: String,
    /// 是否跟随重定向
    #[serde(default)]
    pub follow_redirects: bool,
    /// 最大重定向次数
    #[serde(default)]
    pub max_redirects: u32,
    /// 请求超时时间（秒）
    #[serde(default)]
    pub timeout: u64,
}

/// Nuclei匹配器结构定义
/// 用于匹配HTTP响应中的特征
#[derive(Debug, Deserialize, Clone, Serialize)]
pub struct NucleiMatcher {
    /// 匹配器名称
    #[serde(default)]
    pub name: String,
    /// 状态码匹配列表
    #[serde(default)]
    pub status: Vec<u16>,
    /// 关键词匹配列表
    #[serde(default)]
    pub keyword: Vec<String>,
    /// 正则表达式匹配列表
    #[serde(default)]
    pub regex: Vec<String>,
    /// JSONPath表达式匹配列表
    #[serde(default)]
    pub jsonpath: Vec<String>,
}

/// 提取器结构定义
/// 用于从响应中提取数据
#[derive(Debug, Deserialize, Clone, Serialize)]
pub struct Extractor {
    /// 提取器名称
    #[serde(default)]
    pub name: String,
    /// 提取规则类型
    #[serde(default)]
    pub extractor_type: String,
    /// 提取规则值
    #[serde(default)]
    pub value: String,
}

/// Nuclei信息结构定义
/// 包含模板的基本信息
#[derive(Debug, Deserialize, Clone, Serialize)]
pub struct NucleiInfo {
    /// 模板ID
    #[serde(default)]
    pub id: String,
    /// 模板名称
    #[serde(default)]
    pub name: String,
    /// 模板描述
    #[serde(default)]
    pub description: String,
    /// 模板作者
    #[serde(default)]
    pub author: String,
    /// 严重程度级别
    #[serde(default)]
    pub severity: String,
    /// 参考链接
    #[serde(default)]
    pub reference: Vec<String>,
}

/// Nuclei模板结构定义
/// 表示一个完整的Nuclei扫描模板
#[derive(Debug, Deserialize, Clone, Serialize)]
pub struct NucleiTemplate {
    /// 模板信息
    #[serde(default)]
    pub info: NucleiInfo,
    /// 请求配置
    #[serde(default)]
    pub request: NucleiRequest,
    /// 匹配器列表
    #[serde(default)]
    pub matchers: Vec<NucleiMatcher>,
    /// 提取器列表
    #[serde(default)]
    pub extractors: Vec<Extractor>,
}

/// 响应信息结构定义
/// 包含HTTP响应的相关信息
#[derive(Debug, Clone)]
pub struct ResponseInfo {
    /// 响应状态码
    pub status: u16,
    /// 响应头信息
    pub headers: HashMap<String, String>,
    /// 响应体内容
    pub body: String,
    /// 响应URL
    pub url: String,
    /// 响应时间（毫秒）
    pub response_time: u64,
}

/// Nuclei核心引擎结构定义
/// 提供模板解析、变量替换、匹配器评估等核心功能
#[derive(Debug, Clone)]
pub struct NucleiCoreEngine {
    /// 超时时间配置
    timeout: Duration,
    /// 随机数生成器
    rng: rand::rngs::ThreadRng,
}

impl NucleiCoreEngine {
    /// 创建新的Nuclei核心引擎实例
    /// 
    /// # Returns
    /// - 配置默认值的NucleiCoreEngine实例
    pub fn new() -> Self {
        Self {
            timeout: Duration::from_secs(10),
            rng: rand::thread_rng(),
        }
    }

    /// 创建带有自定义超时时间的Nuclei核心引擎实例
    /// 
    /// # Arguments
    /// - `timeout_seconds`: 超时时间，单位为秒
    /// 
    /// # Returns
    /// - 配置指定超时时间的NucleiCoreEngine实例
    pub fn with_timeout(timeout_seconds: u64) -> Self {
        Self {
            timeout: Duration::from_secs(timeout_seconds),
            rng: rand::thread_rng(),
        }
    }

    /// 从文件加载Nuclei模板
    /// 
    /// # Arguments
    /// - `file_path`: 模板文件路径
    /// 
    /// # Returns
    /// - 成功时返回解析后的NucleiTemplate实例
    pub fn load_template(&self, file_path: &str) -> Result<NucleiTemplate> {
        debug!("Loading template from: {}", file_path);
        
        let content = fs::read_to_string(file_path)?;
        let template: NucleiTemplate = serde_yaml::from_str(&content)?;
        
        debug!("Successfully loaded template: {}", template.info.id);
        Ok(template)
    }

    /// 从目录加载多个Nuclei模板
    /// 
    /// # Arguments
    /// - `dir_path`: 模板目录路径
    /// 
    /// # Returns
    /// - 成功时返回模板路径到模板实例的映射
    pub fn load_templates_from_dir(&self, dir_path: &str) -> Result<HashMap<String, NucleiTemplate>> {
        let path = Path::new(dir_path);
        let mut templates = HashMap::new();
        
        if !path.exists() || !path.is_dir() {
            return Err(anyhow!("Directory does not exist or is not a directory: {}", dir_path));
        }
        
        for entry in fs::read_dir(path)? {
            let entry = entry?;
            let file_path = entry.path();
            
            if file_path.is_file() && file_path.extension().map_or(false, |ext| ext == "yml" || ext == "yaml") {
                match self.load_template(file_path.to_str().unwrap()) {
                    Ok(template) => {
                        templates.insert(file_path.to_str().unwrap().to_string(), template);
                    },
                    Err(e) => {
                        warn!("Failed to load template {}: {}", file_path.display(), e);
                    }
                }
            }
        }
        
        Ok(templates)
    }

    /// 替换字符串中的变量
    /// 
    /// # Arguments
    /// - `input`: 包含变量的输入字符串
    /// - `variables`: 变量映射表
    /// 
    /// # Returns
    /// - 替换后的字符串
    pub fn replace_variables(&mut self, input: &str, variables: &HashMap<String, String>) -> String {
        let mut result = input.to_string();
        let mut replacements = Vec::new();
        
        // 处理 {{randstr:N}} 格式的随机字符串变量
        let randstr_regex = Regex::new(r"\{\{randstr:(\d+)\}\}").unwrap();
        
        for cap in randstr_regex.captures_iter(&result) {
            if let Some(length_str) = cap.get(1) {
                if let Ok(length) = length_str.as_str().parse::<usize>() {
                    let random_str = self.generate_random_string(length);
                    replacements.push((cap[0].to_string(), random_str));
                }
            }
        }
        
        for (pattern, replacement) in replacements {
            result = result.replace(&pattern, &replacement);
        }
        
        // 处理 {{变量名}} 格式的普通变量
        for (var_name, var_value) in variables {
            let pattern = format!("{{{{{}}}}}", var_name);
            result = result.replace(&pattern, var_value);
        }
        
        result
    }

    /// 构建完整的URL
    /// 
    /// # Arguments
    /// - `base_url`: 基础URL
    /// - `path`: 路径部分（可能包含变量）
    /// - `variables`: 变量映射表
    /// 
    /// # Returns
    /// - 完整的URL字符串
    pub fn build_full_url(&mut self, base_url: &str, path: &str, variables: &HashMap<String, String>) -> String {
        // 先替换路径中的变量
        let processed_path = self.replace_variables(path, variables);
        
        // 确保base_url不以/结尾，path不以/开头，然后拼接
        let mut full_url = base_url.to_string();
        if !full_url.ends_with('/') && !processed_path.starts_with('/') {
            full_url.push('/');
        } else if full_url.ends_with('/') && processed_path.starts_with('/') {
            full_url.pop();
        }
        full_url.push_str(&processed_path);
        
        full_url
    }

    /// 评估单个匹配器
    /// 
    /// # Arguments
    /// - `matcher`: 匹配器配置
    /// - `response`: 响应信息
    /// 
    /// # Returns
    /// - 匹配成功返回true，否则返回false
    pub fn evaluate_matcher(&self, matcher: &NucleiMatcher, response: &ResponseInfo) -> bool {
        // 检查状态码匹配
        if !matcher.status.is_empty() && !matcher.status.contains(&response.status) {
            return false;
        }
        
        // 检查关键词匹配
        for keyword in &matcher.keyword {
            if !response.body.contains(keyword) {
                return false;
            }
        }
        
        // 检查正则表达式匹配
        for pattern in &matcher.regex {
            match Regex::new(pattern) {
                Ok(regex) => {
                    if !regex.is_match(&response.body) {
                        return false;
                    }
                },
                Err(e) => {
                    warn!("Invalid regex pattern {}: {}", pattern, e);
                    return false;
                }
            }
        }
        
        // 检查JSONPath匹配
        for jsonpath in &matcher.jsonpath {
            if !self.evaluate_jsonpath(jsonpath, &response.body) {
                return false;
            }
        }
        
        true
    }

    /// 执行所有匹配器
    /// 
    /// # Arguments
    /// - `matchers`: 匹配器列表
    /// - `response`: 响应信息
    /// 
    /// # Returns
    /// - 至少有一个匹配器匹配成功返回true
    pub fn execute_matchers(&self, matchers: &[NucleiMatcher], response: &ResponseInfo) -> bool {
        for matcher in matchers {
            if self.evaluate_matcher(matcher, response) {
                return true;
            }
        }
        false
    }

    /// 评估JSONPath表达式
    /// 
    /// # Arguments
    /// - `jsonpath`: JSONPath表达式
    /// - `body`: 响应体内容
    /// 
    /// # Returns
    /// - 表达式匹配成功返回true，否则返回false
    pub fn evaluate_jsonpath(&self, jsonpath: &str, body: &str) -> bool {
        // 尝试解析响应体为JSON
        match serde_json::from_str::<JsonValue>(body) {
            Ok(json) => {
                // 简化版JSONPath解析
                // 支持 . 分隔的路径和 [] 数组通配符
                let mut current: &JsonValue = &json;
                let parts: Vec<&str> = jsonpath.trim_start_matches('$').split('.').collect();
                
                for part in parts {
                    if part.is_empty() {
                        continue;
                    }
                    
                    // 处理数组通配符
                    if part == "*" {
                        if let Some(array) = current.as_array() {
                            if array.is_empty() {
                                return false;
                            }
                            // 只要有一个元素满足剩余路径，就返回true
                            for item in array {
                                // 简化处理，这里只是检查数组不为空
                                current = item;
                                return true;
                            }
                        }
                        return false;
                    }
                    
                    // 处理普通键
                    current = match current.get(part) {
                        Some(value) => value,
                        None => return false,
                    };
                }
                
                // 路径存在且不为空
                !current.is_null() && current != &JsonValue::Null
            },
            Err(_) => {
                // 响应体不是有效的JSON，直接返回false
                false
            }
        }
    }

    /// 准备变量映射表
    /// 
    /// # Arguments
    /// - `target`: 目标URL
    /// 
    /// # Returns
    /// - 初始化后的变量映射表
    pub fn prepare_variables(&mut self, target: &str) -> HashMap<String, String> {
        let mut variables = HashMap::new();
        
        // 添加基础变量
        variables.insert("target".to_string(), target.to_string());
        variables.insert("baseURL".to_string(), target.to_string());
        
        // 提取域名作为变量
        if let Some(domain) = extract_base_url(target) {
            variables.insert("domain".to_string(), domain);
        }
        
        // 生成随机字符串变量
        variables.insert("randstr".to_string(), self.generate_random_string(10));
        variables.insert("randstr10".to_string(), self.generate_random_string(10));
        variables.insert("randstr20".to_string(), self.generate_random_string(20));
        
        variables
    }

    /// 生成指定长度的随机字符串
    /// 
    /// # Arguments
    /// - `length`: 字符串长度
    /// 
    /// # Returns
    /// - 随机字符串
    pub fn generate_random_string(&mut self, length: usize) -> String {
        const CHARSET: &[u8] = b"ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789";
        
        let mut result = String::with_capacity(length);
        for _ in 0..length {
            let idx = self.rng.gen_range(0..CHARSET.len());
            result.push(CHARSET[idx] as char);
        }
        
        result
    }
}

/// 从URL中提取基础URL部分
/// 
/// # Arguments
/// - `url`: 完整URL
/// 
/// # Returns
/// - 基础URL部分
fn extract_base_url(url: &str) -> Option<String> {
    if let Some(pos) = url.find("://") {
        let scheme_end = pos + 3;
        if let Some(path_start) = url[scheme_end..].find("/") {
            let base_end = scheme_end + path_start;
            return Some(url[0..base_end].to_string());
        } else {
            return Some(url.to_string());
        }
    }
    None
}

#[cfg(test)]
mod tests {
    use super::*;
    
    #[test]
    fn test_load_template() {
        let engine = NucleiCoreEngine::new();
        // 这里可以添加实际的测试用例
        assert!(true);
    }
    
    #[test]
    fn test_replace_variables() {
        let mut engine = NucleiCoreEngine::new();
        let mut variables = HashMap::new();
        variables.insert("target".to_string(), "http://example.com".to_string());
        
        let result = engine.replace_variables("{{target}}/test", &variables);
        assert_eq!(result, "http://example.com/test");
        
        // 测试随机字符串生成
        let result1 = engine.replace_variables("{{randstr:5}}", &HashMap::new());
        let result2 = engine.replace_variables("{{randstr:5}}", &HashMap::new());
        assert_eq!(result1.len(), 5);
        assert_eq!(result2.len(), 5);
    }
    
    #[test]
    fn test_build_full_url() {
        let mut engine = NucleiCoreEngine::new();
        let variables = HashMap::new();
        
        // 测试不同的URL组合情况
        assert_eq!(engine.build_full_url("http://example.com", "path", &variables), "http://example.com/path");
        assert_eq!(engine.build_full_url("http://example.com/", "path", &variables), "http://example.com/path");
        assert_eq!(engine.build_full_url("http://example.com", "/path", &variables), "http://example.com/path");
        assert_eq!(engine.build_full_url("http://example.com/", "/path", &variables), "http://example.com/path");
    }
}